import os,pickle, gensim
import numpy as np
rootDir = os.path.dirname(os.getcwd())
isCodingAndDebug = __name__ == "__main__"

def compute_get_string_vector(string):
    pass

def get_words_matrix():
    matrix_str = np.loadtxt(os.path.join(rootDir, 'GloVe', 'vectors.txt'), dtype=str)
    words,matrix = matrix_str[:,0], matrix_str[:,1:].astype(float)
    word2id = dict(zip(words,range(len(words))))
    return words,word2id,matrix

dictKeyWordsStr2Sentence = pickle.load(open(os.path.join(rootDir,"dictKeyWordsStr2Sentence.pkl"),'rb'))
dictSentence2KeyWordsStr = dict(zip(dictKeyWordsStr2Sentence.values(),dictKeyWordsStr2Sentence.keys()))
keyWordsStrList = dictKeyWordsStr2Sentence.keys()

def get_keyWordsSentence_matrix(word2id,matrix,keyWordsStrList=keyWordsStrList):
    len_words = matrix.shape[0]
    id2keyWordStr = {}
    id2Sentence = {}
    keyWordsStrMatrix = np.zeros((len(keyWordsStrList),len_words),int)
    normalizer = np.zeros((len(keyWordsStrList),1),int)

    func_listWords2listIDs = np.vectorize(lambda word: word2id[word] if word in word2id else word2id["<unk>"])

    def line2count(bundle):
        if isCodingAndDebug : print(bundle)
        ix, line = int(bundle[0]),bundle[1]
        ids = func_listWords2listIDs(line.split())

        def countVecUpdate(idxInWordsList):
            keyWordsStrMatrix[ix][idxInWordsList] += 1
        np.vectorize(countVecUpdate )(ids)

        # for idxInWordsList in ids:
        #     keyWordsStrMatrix[ix][idxInWordsList] += 1
        normalizer[ix]= len(ids)
    bundles = zip(range(len(keyWordsStrList)),keyWordsStrList)
    bundles = list(bundles)
    bundles = np.array(bundles)
    np.apply_along_axis(func1d=line2count,axis=1,arr=bundles)

    unNormalizedMatrix = np.matmul(keyWordsStrMatrix, matrix)
    keyWordsStrMatrix = unNormalizedMatrix / normalizer

    return keyWordsStrMatrix

def sentence2vector(sentence,word2id,matrix):
    keyWordsStr = dictSentence2KeyWordsStr[sentence]
    sentenceVectorMatrix = get_keyWordsSentence_matrix(word2id,matrix,keyWordsStrList=[keyWordsStr])
    if isCodingAndDebug:print(sentenceVectorMatrix)
    return sentenceVectorMatrix

def getTopKids(matrix_q,matrix_v,TopK=5):
    import util
    cosMatrix = util.MatrixCos(matrix_q,matrix_v,isTrimZeros=False)
    sortedIDS = np.argsort(cosMatrix)[:,::-1][:,:TopK]
    if isCodingAndDebug:print(sortedIDS)
    return sortedIDS


if __name__ == "__main__":
    words,word2id,matrix = get_words_matrix()
    keyWordsStrMatrix = get_keyWordsSentence_matrix(word2id, matrix)
    sentencesMatrix1 = sentence2vector(r'(\sin(\alpha))^{n}',word2id,matrix)
    sentencesMatrix2 = sentence2vector(r'\cos{A+B}=\cos{A}*\cos{B}-\sin{A}*\sin{B}', word2id, matrix)
    sentencesMatrix = np.vstack([sentencesMatrix1,sentencesMatrix2])
    getTopKids(sentencesMatrix,keyWordsStrMatrix)
